#include <iostream>
#include "helper_thread_idx.h"
#include "mthcomb.h"
#include "utility.h"
#include "helper_timer.h"
#include "helper_cuda.h"

#define MAX_DSP_COUNT 65535

__global__ void kernel(unsigned char *d_remainderSetAll,
                       unsigned char *d_k1SetAll,
                       unsigned char *d_k2SetAll,
                       unsigned int *d_increment,
                       unsigned char *d_rs,
                       __uint128_t eSetAllCnt,
                       __uint128_t k1_ESetAllCnt,
                       __uint128_t k2_ESetAllCnt,
                       __uint128_t totalCombCnt,
                       __uint128_t batch,
                       __uint128_t iter,
                       int v,
                       int k1,
                       int k2,
                       int e,
                       int lambda)
{
   // global thread index
   int threadId = blockIdx.x * blockDim.x + threadIdx.x;
   // 所取元素的存放首地址
   unsigned char *remainderSetBaseAddr = d_remainderSetAll + threadId * v;
   unsigned char *k1SetBaseAddr = d_k1SetAll + threadId * k1;
   unsigned char *k2SetBaseAddr = d_k2SetAll + threadId * k2;
   for (__uint128_t curIter = 0; curIter < iter; ++curIter)
   {
      __uint128_t cur = curIter * batch + threadId;
      if (cur >= totalCombCnt)
         return; // end of combination

      // 重置 remainderSet
      for (int i = 0; i < v; ++i)
      {
         remainderSetBaseAddr[i] = 1;
      }

      // 取集合eSet
      if (e != 0)
      {
         __uint128_t eSetIdx = (cur / (k1_ESetAllCnt * k2_ESetAllCnt)) % eSetAllCnt;
         mThComb(eSetIdx, eSetAllCnt, v, e, k1SetBaseAddr);

         // 组合元素即为所取元素
         for (int i = 0; i < e; ++i)
         {
            remainderSetBaseAddr[k1SetBaseAddr[i]] = 0; // 对应元素已被取走
            k2SetBaseAddr[i] = k1SetBaseAddr[i];
         }
      }

      // 取集合k1-eSet
      if ((k1 - e) != 0)
      {
         __uint128_t k1_eSetIdx = (cur / k2_ESetAllCnt) % k1_ESetAllCnt;
         mThComb(k1_eSetIdx, k1_ESetAllCnt, v - e, k1 - e, k1SetBaseAddr + e);

         int tempIdx = 0;
         int orderCnt = -1;
         for (int i = 0; i < v; ++i)
         {
            if (remainderSetBaseAddr[i] == 1)
            {
               ++orderCnt;
            }
            if (orderCnt == k1SetBaseAddr[e + tempIdx])
            {
               // 当前元素应取
               k1SetBaseAddr[e + tempIdx] = i;
               remainderSetBaseAddr[i] = 0;
               ++tempIdx;
            }
         }
      }

      // 取集合k2-eSet
      if ((k2 - e) != 0)
      {
         __uint128_t k2_eSetIdx = cur % k2_ESetAllCnt;
         mThComb(k2_eSetIdx, k2_ESetAllCnt, v - k1, k2 - e, k2SetBaseAddr + e);

         int tempIdx = 0;
         int orderCnt = -1;
         for (int i = 0; i < v; ++i)
         {
            if (remainderSetBaseAddr[i] == 1)
            {
               ++orderCnt;
            }
            if (orderCnt == k2SetBaseAddr[e + tempIdx])
            {
               // 当前元素应取
               k2SetBaseAddr[e + tempIdx] = i;
               remainderSetBaseAddr[i] = 0;
               ++tempIdx;
            }
         }
      }

      // 重用remainderSetBaseAddr地址，作为差表使用
      for (int i = 0; i < v; ++i)
      {
         remainderSetBaseAddr[i] = 0;
      }

      // 计算差表
      for (int i = 0; i < k1; ++i)
      {
         for (int j = 0; j < k2; ++j)
         {
            remainderSetBaseAddr[(k1SetBaseAddr[i] - k2SetBaseAddr[j] + v) % v] += 1;
         }
      }

      // 判断差集偶
      bool dspFlag = true;
      for (int i = 1; i < v; ++i) // 从1开始判断
      {
         if (remainderSetBaseAddr[i] != lambda)
         {
            dspFlag = false;
            break;
         }
      }

      // 利用原子加保存结果
      if (dspFlag)
      {
         // 原子加 读取位于全局或共享存储器中地址address 处的32 位字old，计算 ((old >= val) ? 0 : (old+1))，
         // 并将结果存储在存储器的同一地址中。这三项操作在一次原子事务中执行。该函数将返回old。
         unsigned int rsIdx = atomicInc(d_increment, MAX_DSP_COUNT - 1);
         unsigned char *d_rsBaseAddr = d_rs + rsIdx * (k1 + k2);
         // 存储结果
         for (int i = 0; i < k1; ++i)
         {
            d_rsBaseAddr[i] = k1SetBaseAddr[i];
         }

         for (int j = 0; j < k2; ++j)
         {
            d_rsBaseAddr[k1 + j] = k2SetBaseAddr[j];
         }
      }
   }
}

void print_uint128_t(__uint128_t value)
{
   unsigned long long upper = value >> 64;
   unsigned long long lower = (unsigned long long)value;
   printf("%llu%016llu", upper, lower);
}

void save(int v, int k1, int k2, int e, int lambda, unsigned char *rs, unsigned int *increment)
{
   char buffer[255];
   sprintf(buffer, "dsp-%d-%d-%d-%d-%d.txt", v, k1, k2, e, lambda); // concatenated file name, eg: dsp-15-3-5-1-1.txt

   FILE *fp = fopen(buffer, "w");
   if (fp != NULL)
   {
      for (size_t i = 0; i < *increment; ++i)
      {
         unsigned char *rsBaseAddr = rs + i * (k1 + k2);
         // 打印输出dsp
         fprintf(fp, "[[%d", rsBaseAddr[0]);
         for (int j = 1; j < k1; ++j)
         {
            fprintf(fp, ",%d", rsBaseAddr[j]);
         }
         fprintf(fp, "],[%d", rsBaseAddr[k1]);
         for (int j = k1 + 1; j < (k1 + k2); ++j)
         {
            fprintf(fp, ",%d", rsBaseAddr[j]);
         }
         fprintf(fp, "]]\n");
      }
      fclose(fp); // release resource
   }
}

void doSearch(int v, int k1, int k2, int e, int lambda)
{
   // prepare profiling
   StopWatchInterface *timer = nullptr;
   sdkCreateTimer(&timer);
   sdkResetTimer(&timer);

   __uint128_t eSetAllCnt = e == 0 ? 0 : choose(v, e);
   __uint128_t k1_ESetAllCnt = k1 - e == 0 ? 0 : choose(v - e, k1 - e);
   __uint128_t k2_ESetAllCnt = k2 - e == 0 ? 0 : choose(v - k1, k2 - e);
   __uint128_t totalCombCnt = 1;
   totalCombCnt *= eSetAllCnt == 0 ? 1 : eSetAllCnt;
   totalCombCnt *= k1_ESetAllCnt == 0 ? 1 : k1_ESetAllCnt;
   totalCombCnt *= k2_ESetAllCnt == 0 ? 1 : k2_ESetAllCnt;

   __uint128_t blocks = 1024;
   __uint128_t threads = 512;
   __uint128_t batch = blocks * threads; // 一个批量的线程数

   unsigned char *d_remainderSetAll = 0; // 用于存储剩余元素 根据组合元素来选取剩余元素中的元素
   unsigned char *d_k1SetAll = 0;        // 用于存储k1Set
   unsigned char *d_k2SetAll = 0;        // 用于存储k2Set

   checkCudaErrors(cudaMalloc((void **)&d_remainderSetAll, batch * v * sizeof(unsigned char)));
   checkCudaErrors(cudaMalloc((void **)&d_k1SetAll, batch * k1 * sizeof(unsigned char)));
   checkCudaErrors(cudaMalloc((void **)&d_k2SetAll, batch * k2 * sizeof(unsigned char)));

   // 利用原子加获取线程存放结果的索引位置
   unsigned int *d_increment = 0; // 原子加返回该变量存放的值，同时自增1
   unsigned char *d_rs = 0;       // 结果存放位置

   checkCudaErrors(cudaMalloc((void **)&d_increment, sizeof(unsigned int)));
   checkCudaErrors(cudaMalloc((void **)&d_rs, MAX_DSP_COUNT * (k1 + k2) * sizeof(unsigned char)));

   checkCudaErrors(cudaMemset(d_increment, 0, sizeof(unsigned int)));
   checkCudaErrors(cudaMemset(d_rs, 0, MAX_DSP_COUNT * (k1 + k2) * sizeof(unsigned char)));

   __uint128_t iter = totalCombCnt % batch == 0 ? (totalCombCnt / batch) : (totalCombCnt / batch + 1);
   printf("The process would repeat ");
   print_uint128_t(iter);
   printf(" times.\n");
   printf("CUDA kernel launched with ");
   print_uint128_t(blocks);
   printf(" blocks of ");
   print_uint128_t(threads);
   printf(" threads.\n");
   printf("Start timer.\n");

   // start profiling
   sdkStartTimer(&timer);
   kernel<<<blocks, threads>>>(d_remainderSetAll,
                               d_k1SetAll,
                               d_k2SetAll,
                               d_increment,
                               d_rs,
                               eSetAllCnt,
                               k1_ESetAllCnt,
                               k2_ESetAllCnt,
                               totalCombCnt,
                               batch,
                               iter,
                               v,
                               k1,
                               k2,
                               e,
                               lambda);

   checkCudaErrors(cudaDeviceSynchronize()); // wait for device to complete
   cudaError_t err = cudaGetLastError();
   if (err != cudaSuccess)
   {
      fprintf(stderr, "CUDA error at %s:%d \"%s\" \"%s\" \n", __FILE__, __LINE__, __func__, cudaGetErrorString(err));
   }

   sdkStopTimer(&timer);

   unsigned char *rs = new unsigned char[MAX_DSP_COUNT * (k1 + k2)];
   unsigned int *increment = new unsigned int;
   checkCudaErrors(cudaMemcpy(rs, d_rs, MAX_DSP_COUNT * (k1 + k2) * sizeof(unsigned char), cudaMemcpyDeviceToHost));
   checkCudaErrors(cudaMemcpy(increment, d_increment, sizeof(unsigned int), cudaMemcpyDeviceToHost));

   // 保存结果
   save(v, k1, k2, e, lambda, rs, increment);
   printf("dsp count: %d\n\n\n", *increment);
   printf("time spent : %.2f ms\n", sdkGetTimerValue(&timer));

   // 释放内存
   checkCudaErrors(cudaFree(d_remainderSetAll));
   checkCudaErrors(cudaFree(d_k1SetAll));
   checkCudaErrors(cudaFree(d_k2SetAll));
   checkCudaErrors(cudaFree(d_increment));
   checkCudaErrors(cudaFree(d_rs));

   delete[] rs;
   delete increment;
}

int main(int argc, char *argv[])
{
   greeting(argc, argv);

   int v, k1, k2, e, lambda;
   printf("please input v k1 k2 e lambda:\n");
   scanf("%d %d %d %d %d", &v, &k1, &k2, &e, &lambda);
#ifndef NDEBUG
   printf("you entered %d %d %d %d %d\n", v, k1, k2, e, lambda);
#endif

   if (k1 * k2 != lambda * (v - 1) + e)
   {
      printf("invalid input!\n");
      return -1;
   }

   printf("loading, please wait...\n");
   doSearch(v, k1, k2, e, lambda);
}